# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Neighbour list
"""

from inspect import signature
from distutils.version import LooseVersion
from typing import Tuple
import torch
from torch import Tensor
from torch.nn import Parameter, Module

from . import FullConnectNeighbours, DistanceNeighbours, GridNeighbours
from ..system import Molecule
from ..function.functions import gather_vector, get_integer, get_tensor
from ..function.operations import GetVector


class NeighbourList(Module):
    r"""Neighbour list

    Args:

        system (Molecule):      Simulation system.

        cutoff (float):         Cut-off distance. If None is given under periodic boundary condition (PBC),
                                the cutoff will be assigned with the default value of 1 nm.
                                Default: None

        pace (int):             Update frequency for neighbour list. Default: 20

        exclude_index (Tensor): Tensor of the indices of the neighbouring atoms which could be excluded from the
                                neighbour list. The shape of Tensor is `(B, A, Ex)`, and the data type is int.
                                Default: None

        num_neighbours (int):   Maximum number of neighbours. If `None` is given, this value will be calculated
                                by the ratio of the number of neighbouring grids to the total number of grids.
                                Default: None

        num_cell_cut (int):     Number of subdivision of grid cells according to cutoff. Default: 1

        cutoff_scale (float):   Factor to scale cutoff distance. Default: 1.2

        cell_cap_scale (float): Scale factor for `cell_capacity`. Default: 1.25

        grid_num_scale (float): Scale factor to calculate `num_neighbours` by ratio of grids.
                                If `num_neighbours` is not None, it will not be used. Default: 1.5

        large_dis (float):      A large number to fill in the distances to the masked neighbouring atoms.
                                Default: 1e4

        use_grids (bool):       Whether to use grids to calculate the neighbour list. Default: None

        cast_fp16 (bool):       If this is set to `True`, the data will be cast to float16 before sort.
                                For use with some devices that only support sorting of float16 data.
                                Default: False

    Supported Platforms:

        ``Ascend`` ``GPU``

    Symbols:

        B:  Batchsize, i.e. number of walkers of the simulation.

        A:  Number of the atoms in the simulation system.

        N:  Number of the maximum neighbouring atoms.

        D:  Dimension of position coordinates.

        Ex: Maximum number of excluded neighbour atoms.

    """

    def __init__(self,
                 system: Molecule,
                 cutoff: float = None,
                 pace: int = 20,
                 exclude_index: Tensor = None,
                 num_neighbours: int = None,
                 num_cell_cut: int = 1,
                 cutoff_scale: float = 1.05,
                 cell_cap_scale: float = 1.05,
                 grid_num_scale: float = 2,
                 use_grids: bool = False,
                 centroid: bool = False,
                 ):

        super().__init__()


        self.num_walker = system.num_walker
        self.coordinate = system.coordinate if not centroid else system.coordinate.mean(dim=0, keepdim=True)
        self.pbc_box = system.pbc_box.mean(dim=0, keepdim=True) if centroid and system.pbc_box is not None else system.pbc_box
        self.num_atoms = self.coordinate.shape[-2]
        self.dim = self.coordinate.shape[-1]
        use_pbc = self.pbc_box is not None

        self.atom_mask = system.atom_mask
        self.exclude_index = exclude_index
        if exclude_index is not None:
            self.exclude_index = get_tensor(exclude_index, dtype=torch.int32)

        self.use_grids = use_grids

        self._pace = get_integer(pace)
        if self._pace < 0:
            raise ValueError('pace cannot be less than 0!')

        if cutoff is None and self.pbc_box is not None:
            cutoff = self.pbc_box.min() / 2
            print(f'[WARNING] The cutoff distance is not given. '
                  f'It will be assigned with the half of the minimum of the box, which is {cutoff}')

        self.no_mask = False
        if cutoff is None:
            self.cutoff = None
            self.large_dis = get_tensor(1e4, dtype=torch.float32)
            self._pace = 0
            self.neighbour_list = FullConnectNeighbours(self.num_atoms)
            if self.exclude_index is None:
                self.no_mask = True
        else:
            self.cutoff = get_tensor(cutoff, dtype=torch.float32)
            self.large_dis = self.cutoff * 100
            if self.use_grids or self.use_grids is None:
                self.neighbour_list = GridNeighbours(
                    cutoff=self.cutoff,
                    coordinate=self.coordinate,
                    pbc_box=self.pbc_box,
                    atom_mask=self.atom_mask,
                    exclude_index=self.exclude_index,
                    num_neighbours=num_neighbours,
                    num_cell_cut=num_cell_cut,
                    cutoff_scale=cutoff_scale,
                    cell_cap_scale=cell_cap_scale,
                    grid_num_scale=grid_num_scale,
                )
                if self.neighbour_list.neigh_capacity >= self.num_atoms:
                    if self.use_grids is True:
                        print(f'[WARNING] The number of neighbour atoms in `GridNeighbours` '
                              f'({self.neighbour_list.neigh_capacity}) is not less than '
                              f'the number of atoms ({self.num_atoms}). '
                              f'It would be more efficient to use `DistanceNeighbours` '
                              f'(set `use_grids` to `False` or `None`).')
                    else:
                        self.use_grids = False
                else:
                    self.use_grids = True

            if not self.use_grids:
                self.neighbour_list = DistanceNeighbours(
                    cutoff=self.cutoff,
                    num_neighbours=num_neighbours,
                    atom_mask=self.atom_mask,
                    exclude_index=self.exclude_index,
                    use_pbc=use_pbc,
                    cutoff_scale=cutoff_scale,
                    large_dis=self.large_dis,
                )

        self.num_neighbours = num_neighbours

        self.neighbours = None
        self.neighbour_mask = None

        if isinstance(self.neighbour_list, FullConnectNeighbours):
            index, mask = self.neighbour_list(self.atom_mask, self.exclude_index)
        else:
            index, mask = self.neighbour_list(self.coordinate, self.pbc_box, self.atom_mask, self.exclude_index)

        if index is not None:
            self.neighbours = index
            if mask is not None:
                self.neighbour_mask = mask

        self.get_vector = GetVector(use_pbc)

    @property
    def pace(self) -> int:
        """Get update frequency"""
        return self._pace

    def set_exclude_index(self, exclude_index: Tensor):
        """Set exclude index"""
        self.exclude_index = exclude_index
        self.neighbour_list.set_exclude_index(exclude_index)
        index, mask = self.update(self.coordinate, self.pbc_box)
        self.neighbours = index
        self.neighbour_mask = mask
        return self

    def print_info(self):
        """Print neighbor list information and return as a dictionary"""
        info = {}
        
        # Basic information
        info['Update Frequency'] = self._pace
        info['Periodic Boundary'] = self.pbc_box is not None
        info['Number of Atoms'] = self.num_atoms
        
        if isinstance(self.neighbour_list, FullConnectNeighbours):
            info['NeighbourList Type'] = 'Full Connect'
        else:
            info['NeighbourList Type'] = 'Grid' if isinstance(self.neighbour_list, GridNeighbours) else 'Distance'
            info['Cutoff'] = float(self.cutoff)
            info['Max Neighbors'] = self.num_neighbours if self.num_neighbours else 'auto'
            
            if isinstance(self.neighbour_list, GridNeighbours):
                info['Grid Cell_cuts'] = self.neighbour_list.num_cell_cut
                info['Cell Capacity Scale'] = float(self.neighbour_list.cell_cap_scale)
            elif isinstance(self.neighbour_list, DistanceNeighbours):
                info['Cutoff Scale'] = float(self.neighbour_list.cutoff_scale)
        
        if self.exclude_index is not None:
            info['Number of Excluded'] = self.exclude_index.shape[1]
            
        return info

    def update(self, coordinate: Tensor, pbc_box: Tensor = None) -> Tuple[Tensor, Tensor]:
        """Update neighbour list"""
        if self.pace == 0:
            return self.neighbours, self.neighbour_mask

        with torch.no_grad():
            index, mask = self.calculate(coordinate, pbc_box)

        if index is not None:
            self.neighbours = index
            if mask is not None:
                self.neighbour_mask = mask

        return index, mask

    def calculate(self, coordinate: Tensor, pbc_box: Tensor = None) -> Tuple[Tensor, Tensor]:
        """Calculate neighbour list"""
        if self.pace == 0:
            return self.neighbours, self.neighbour_mask
        
        if self.cutoff is None:
            return self.neighbour_list(self.atom_mask, self.exclude_index)

        index, mask = self.neighbour_list(coordinate, pbc_box, self.atom_mask, self.exclude_index)

        return index, mask

    def get_neighbour_list(self) -> Tuple[Tensor, Tensor]:
        """Get neighbour list"""
        if self.neighbours is None:
            return None, None
        return self.neighbours, self.neighbour_mask

    def forward(self,
                  coordinate: Tensor,
                  pbc_box: Tensor = None
                  ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
        """Forward"""

        index, mask = self.get_neighbour_list()
        # (B, A, N, D)
        vectors = gather_vector(coordinate, index)
        vectors = self.get_vector(coordinate.unsqueeze(-2), vectors, pbc_box)

        if mask is not None:
            large_dis = torch.broadcast_to(self.large_dis, mask.shape)
            large_dis = torch.where(mask, torch.zeros_like(large_dis), large_dis)
            # (B, A, N, D) = (B, A, N, D) + (B, A, N, 1)
            vectors = vectors + large_dis.unsqueeze(-1)
        # (B, A, N)
        distances = torch.norm(vectors, dim=-1)

        return index, mask, vectors, distances
